#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
增强版可视化模块

根据优化方案实现的可视化稳定性增强：
1. 字体管理器 - 解决中文显示问题
2. 内存管理器 - 防止内存泄漏
3. 异常处理机制 - 提高稳定性
4. 图表生成器 - 统一接口
"""

import os
import gc
import platform
import logging
import warnings
from typing import Dict, List, Any, Optional, Tuple, Union
from dataclasses import dataclass
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.font_manager as fm
from matplotlib.font_manager import FontProperties
import seaborn as sns
from contextlib import contextmanager
import psutil
from pathlib import Path

# 设置matplotlib后端
matplotlib.use('Agg')  # 使用非交互式后端

# 忽略matplotlib警告
warnings.filterwarnings('ignore', category=UserWarning, module='matplotlib')


@dataclass
class VisualizationConfig:
    """可视化配置"""
    figure_size: Tuple[int, int] = (10, 6)
    dpi: int = 100
    font_size: int = 12
    title_font_size: int = 14
    label_font_size: int = 10
    color_palette: str = 'Set2'
    style: str = 'whitegrid'
    max_memory_mb: int = 500
    max_figures: int = 10
    auto_cleanup: bool = True


class FontManager:
    """
    字体管理器 - 解决中文显示问题
    """
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.current_font = None
        self.font_cache = {}
        
        # 系统字体路径配置
        self.font_paths = {
            'windows': [
                'C:/Windows/Fonts/simhei.ttf',  # 黑体
                'C:/Windows/Fonts/msyh.ttf',    # 微软雅黑
                'C:/Windows/Fonts/simsun.ttc',  # 宋体
                'C:/Windows/Fonts/simkai.ttf'   # 楷体
            ],
            'darwin': [  # macOS
                '/System/Library/Fonts/PingFang.ttc',
                '/System/Library/Fonts/STHeiti Light.ttc',
                '/System/Library/Fonts/Hiragino Sans GB.ttc'
            ],
            'linux': [
                '/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf',
                '/usr/share/fonts/truetype/wqy/wqy-microhei.ttc',
                '/usr/share/fonts/truetype/arphic/ukai.ttc'
            ]
        }
        
        # 自动配置字体
        self._auto_configure_font()
    
    def _auto_configure_font(self) -> bool:
        """
        自动检测并配置中文字体
        """
        try:
            system = platform.system().lower()
            
            # 根据操作系统选择字体候选列表
            if 'windows' in system:
                font_candidates = self.font_paths['windows']
            elif 'darwin' in system:
                font_candidates = self.font_paths['darwin']
            else:
                font_candidates = self.font_paths['linux']
            
            # 尝试配置字体
            for font_path in font_candidates:
                if self._try_configure_font(font_path):
                    self.logger.info(f"成功配置中文字体: {font_path}")
                    return True
            
            # 如果没有找到合适的字体，使用备用方案
            self._configure_fallback_font()
            self.logger.warning("使用备用字体配置")
            return False
            
        except Exception as e:
            self.logger.error(f"字体配置失败: {e}")
            self._configure_fallback_font()
            return False
    
    def _try_configure_font(self, font_path: str) -> bool:
        """
        尝试配置指定字体
        """
        try:
            if os.path.exists(font_path):
                # 测试字体是否可用
                font_prop = FontProperties(fname=font_path)
                
                # 配置matplotlib
                plt.rcParams['font.sans-serif'] = [font_prop.get_name()]
                plt.rcParams['axes.unicode_minus'] = False
                plt.rcParams['font.size'] = 12
                
                # 缓存字体属性
                self.current_font = font_path
                self.font_cache['default'] = font_prop
                
                return True
        except Exception as e:
            self.logger.debug(f"字体配置失败 {font_path}: {e}")
        
        return False
    
    def _configure_fallback_font(self):
        """
        配置备用字体
        """
        try:
            # 使用matplotlib默认字体
            plt.rcParams['font.sans-serif'] = ['DejaVu Sans', 'Arial', 'sans-serif']
            plt.rcParams['axes.unicode_minus'] = False
            
            # 尝试使用系统中文字体
            chinese_fonts = ['SimHei', 'Microsoft YaHei', 'PingFang SC', 'WenQuanYi Micro Hei']
            available_fonts = [f.name for f in fm.fontManager.ttflist]
            
            for font in chinese_fonts:
                if font in available_fonts:
                    plt.rcParams['font.sans-serif'].insert(0, font)
                    break
            
            self.current_font = 'fallback'
            
        except Exception as e:
            self.logger.error(f"备用字体配置失败: {e}")
    
    def get_font_properties(self, size: Optional[int] = None) -> FontProperties:
        """
        获取字体属性
        
        Args:
            size: 字体大小
            
        Returns:
            FontProperties: 字体属性对象
        """
        if self.current_font and self.current_font != 'fallback':
            font_prop = FontProperties(fname=self.current_font)
            if size:
                font_prop.set_size(size)
            return font_prop
        else:
            font_prop = FontProperties()
            if size:
                font_prop.set_size(size)
            return font_prop
    
    def test_chinese_display(self) -> bool:
        """
        测试中文显示是否正常
        
        Returns:
            bool: 是否支持中文显示
        """
        try:
            fig, ax = plt.subplots(figsize=(6, 4))
            ax.text(0.5, 0.5, '中文测试', fontproperties=self.get_font_properties(),
                   ha='center', va='center', fontsize=16)
            ax.set_title('字体测试', fontproperties=self.get_font_properties())
            plt.close(fig)
            return True
        except Exception as e:
            self.logger.error(f"中文显示测试失败: {e}")
            return False


class MemoryManager:
    """
    内存管理器 - 防止内存泄漏
    """
    
    def __init__(self, max_memory_mb: int = 500, max_figures: int = 10):
        self.logger = logging.getLogger(__name__)
        self.max_memory_mb = max_memory_mb
        self.max_figures = max_figures
        self.active_figures = []
        self.memory_usage_history = []
        
        # 获取进程对象
        self.process = psutil.Process()
    
    def get_memory_usage(self) -> float:
        """
        获取当前内存使用量（MB）
        """
        try:
            memory_info = self.process.memory_info()
            return memory_info.rss / 1024 / 1024  # 转换为MB
        except Exception:
            return 0.0
    
    def check_memory_limit(self) -> bool:
        """
        检查是否超过内存限制
        """
        current_memory = self.get_memory_usage()
        self.memory_usage_history.append(current_memory)
        
        # 保持历史记录长度
        if len(self.memory_usage_history) > 100:
            self.memory_usage_history = self.memory_usage_history[-50:]
        
        return current_memory > self.max_memory_mb
    
    @contextmanager
    def managed_figure(self, figsize: Tuple[int, int] = (10, 6), **kwargs):
        """
        管理图形的上下文管理器
        
        Args:
            figsize: 图形大小
            **kwargs: 其他matplotlib参数
        """
        fig = None
        try:
            # 检查内存使用
            if self.check_memory_limit():
                self.cleanup_old_figures(force=True)
                gc.collect()
            
            # 检查图形数量限制
            if len(self.active_figures) >= self.max_figures:
                self.cleanup_old_figures(count=len(self.active_figures) - self.max_figures + 1)
            
            # 创建新图形
            fig = plt.figure(figsize=figsize, **kwargs)
            self.active_figures.append(fig)
            
            yield fig
            
        except Exception as e:
            self.logger.error(f"图形创建失败: {e}")
            if fig:
                plt.close(fig)
            raise
        finally:
            # 清理资源
            if fig and fig in self.active_figures:
                self.active_figures.remove(fig)
    
    def cleanup_old_figures(self, count: Optional[int] = None, force: bool = False):
        """
        清理旧图形
        
        Args:
            count: 清理数量，None表示清理一半
            force: 是否强制清理所有
        """
        try:
            if force:
                # 强制清理所有图形
                for fig in self.active_figures:
                    plt.close(fig)
                self.active_figures.clear()
            else:
                # 清理指定数量的图形
                cleanup_count = count or max(1, len(self.active_figures) // 2)
                for _ in range(min(cleanup_count, len(self.active_figures))):
                    old_fig = self.active_figures.pop(0)
                    plt.close(old_fig)
            
            # 强制垃圾回收
            gc.collect()
            
            self.logger.debug(f"清理了图形，当前活跃图形数: {len(self.active_figures)}")
            
        except Exception as e:
            self.logger.error(f"图形清理失败: {e}")
    
    def get_memory_stats(self) -> Dict[str, Any]:
        """
        获取内存统计信息
        """
        current_memory = self.get_memory_usage()
        
        stats = {
            'current_memory_mb': current_memory,
            'max_memory_mb': self.max_memory_mb,
            'memory_usage_ratio': current_memory / self.max_memory_mb,
            'active_figures': len(self.active_figures),
            'max_figures': self.max_figures
        }
        
        if self.memory_usage_history:
            stats.update({
                'avg_memory_mb': np.mean(self.memory_usage_history),
                'peak_memory_mb': np.max(self.memory_usage_history),
                'memory_trend': 'increasing' if len(self.memory_usage_history) > 1 and 
                               self.memory_usage_history[-1] > self.memory_usage_history[-2] else 'stable'
            })
        
        return stats


class EnhancedVisualizationEngine:
    """
    增强版可视化引擎
    """
    
    def __init__(self, config: Optional[VisualizationConfig] = None):
        self.config = config or VisualizationConfig()
        self.logger = logging.getLogger(__name__)
        
        # 初始化组件
        self.font_manager = FontManager()
        self.memory_manager = MemoryManager(
            max_memory_mb=self.config.max_memory_mb,
            max_figures=self.config.max_figures
        )
        
        # 配置seaborn样式
        self._configure_style()
        
        # 测试中文显示
        if not self.font_manager.test_chinese_display():
            self.logger.warning("中文字体可能无法正常显示")
    
    def _configure_style(self):
        """
        配置可视化样式
        """
        try:
            # 设置seaborn样式
            sns.set_style(self.config.style)
            sns.set_palette(self.config.color_palette)
            
            # 设置matplotlib参数
            plt.rcParams.update({
                'figure.figsize': self.config.figure_size,
                'figure.dpi': self.config.dpi,
                'font.size': self.config.font_size,
                'axes.titlesize': self.config.title_font_size,
                'axes.labelsize': self.config.label_font_size,
                'xtick.labelsize': self.config.label_font_size,
                'ytick.labelsize': self.config.label_font_size,
                'legend.fontsize': self.config.label_font_size,
                'figure.autolayout': True,
                'savefig.bbox': 'tight',
                'savefig.dpi': self.config.dpi
            })
            
        except Exception as e:
            self.logger.error(f"样式配置失败: {e}")
    
    def create_scenario_comparison_chart(self, scenario_data: Dict[str, Any], 
                                       save_path: Optional[str] = None) -> str:
        """
        创建情景对比图表
        
        Args:
            scenario_data: 情景数据
            save_path: 保存路径
            
        Returns:
            str: 图表保存路径或base64编码
        """
        try:
            with self.memory_manager.managed_figure(figsize=(12, 8)) as fig:
                # 创建子图
                gs = fig.add_gridspec(2, 2, hspace=0.3, wspace=0.3)
                
                # 情景利润对比
                ax1 = fig.add_subplot(gs[0, :])
                scenarios = list(scenario_data.keys())
                profits = [scenario_data[s].get('profit', 0) for s in scenarios]
                
                bars = ax1.bar(scenarios, profits, color=sns.color_palette(self.config.color_palette, len(scenarios)))
                ax1.set_title('各情景利润对比', fontproperties=self.font_manager.get_font_properties(self.config.title_font_size))
                ax1.set_ylabel('利润', fontproperties=self.font_manager.get_font_properties())
                
                # 添加数值标签
                for bar, profit in zip(bars, profits):
                    height = bar.get_height()
                    ax1.text(bar.get_x() + bar.get_width()/2., height + height*0.01,
                            f'{profit:.2f}', ha='center', va='bottom',
                            fontproperties=self.font_manager.get_font_properties())
                
                # 决策变量热力图
                ax2 = fig.add_subplot(gs[1, 0])
                decision_matrix = []
                variable_names = []
                
                for scenario, data in scenario_data.items():
                    decisions = data.get('decisions', {})
                    if not variable_names:
                        variable_names = list(decisions.keys())
                    decision_matrix.append([decisions.get(var, 0) for var in variable_names])
                
                if decision_matrix:
                    im = ax2.imshow(decision_matrix, cmap='RdYlBu_r', aspect='auto')
                    ax2.set_xticks(range(len(variable_names)))
                    ax2.set_xticklabels(variable_names, fontproperties=self.font_manager.get_font_properties())
                    ax2.set_yticks(range(len(scenarios)))
                    ax2.set_yticklabels(scenarios, fontproperties=self.font_manager.get_font_properties())
                    ax2.set_title('决策变量热力图', fontproperties=self.font_manager.get_font_properties())
                    
                    # 添加颜色条
                    plt.colorbar(im, ax=ax2)
                
                # 利润分布箱线图
                ax3 = fig.add_subplot(gs[1, 1])
                ax3.boxplot(profits, labels=['利润分布'])
                ax3.set_title('利润分布', fontproperties=self.font_manager.get_font_properties())
                ax3.set_ylabel('利润', fontproperties=self.font_manager.get_font_properties())
                
                # 保存图表
                if save_path:
                    fig.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
                    return save_path
                else:
                    return self._fig_to_base64(fig)
                    
        except Exception as e:
            self.logger.error(f"情景对比图表创建失败: {e}")
            return self._create_error_chart(f"图表创建失败: {str(e)}")
    
    def create_optimization_convergence_chart(self, convergence_data: Dict[str, List[float]], 
                                            save_path: Optional[str] = None) -> str:
        """
        创建优化收敛图表
        
        Args:
            convergence_data: 收敛数据 {algorithm: [fitness_values]}
            save_path: 保存路径
            
        Returns:
            str: 图表保存路径或base64编码
        """
        try:
            with self.memory_manager.managed_figure(figsize=(10, 6)) as fig:
                ax = fig.add_subplot(111)
                
                colors = sns.color_palette(self.config.color_palette, len(convergence_data))
                
                for i, (algorithm, values) in enumerate(convergence_data.items()):
                    iterations = range(1, len(values) + 1)
                    ax.plot(iterations, values, label=algorithm, color=colors[i], linewidth=2)
                
                ax.set_xlabel('迭代次数', fontproperties=self.font_manager.get_font_properties())
                ax.set_ylabel('目标函数值', fontproperties=self.font_manager.get_font_properties())
                ax.set_title('算法收敛性对比', fontproperties=self.font_manager.get_font_properties(self.config.title_font_size))
                ax.legend(prop=self.font_manager.get_font_properties())
                ax.grid(True, alpha=0.3)
                
                # 保存图表
                if save_path:
                    fig.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
                    return save_path
                else:
                    return self._fig_to_base64(fig)
                    
        except Exception as e:
            self.logger.error(f"收敛图表创建失败: {e}")
            return self._create_error_chart(f"图表创建失败: {str(e)}")
    
    def create_pareto_frontier_chart(self, pareto_data: List[Dict[str, float]], 
                                   objective_names: List[str],
                                   save_path: Optional[str] = None) -> str:
        """
        创建帕累托前沿图表
        
        Args:
            pareto_data: 帕累托解集数据
            objective_names: 目标函数名称
            save_path: 保存路径
            
        Returns:
            str: 图表保存路径或base64编码
        """
        try:
            with self.memory_manager.managed_figure(figsize=(10, 8)) as fig:
                if len(objective_names) == 2:
                    # 二维帕累托前沿
                    ax = fig.add_subplot(111)
                    
                    x_values = [point[objective_names[0]] for point in pareto_data]
                    y_values = [point[objective_names[1]] for point in pareto_data]
                    
                    ax.scatter(x_values, y_values, c='red', s=50, alpha=0.7, label='帕累托最优解')
                    ax.plot(x_values, y_values, 'r--', alpha=0.5)
                    
                    ax.set_xlabel(objective_names[0], fontproperties=self.font_manager.get_font_properties())
                    ax.set_ylabel(objective_names[1], fontproperties=self.font_manager.get_font_properties())
                    ax.set_title('帕累托前沿', fontproperties=self.font_manager.get_font_properties(self.config.title_font_size))
                    ax.legend(prop=self.font_manager.get_font_properties())
                    ax.grid(True, alpha=0.3)
                    
                elif len(objective_names) == 3:
                    # 三维帕累托前沿
                    ax = fig.add_subplot(111, projection='3d')
                    
                    x_values = [point[objective_names[0]] for point in pareto_data]
                    y_values = [point[objective_names[1]] for point in pareto_data]
                    z_values = [point[objective_names[2]] for point in pareto_data]
                    
                    ax.scatter(x_values, y_values, z_values, c='red', s=50, alpha=0.7)
                    
                    ax.set_xlabel(objective_names[0], fontproperties=self.font_manager.get_font_properties())
                    ax.set_ylabel(objective_names[1], fontproperties=self.font_manager.get_font_properties())
                    ax.set_zlabel(objective_names[2], fontproperties=self.font_manager.get_font_properties())
                    ax.set_title('三维帕累托前沿', fontproperties=self.font_manager.get_font_properties(self.config.title_font_size))
                    
                else:
                    # 多维目标的平行坐标图
                    ax = fig.add_subplot(111)
                    
                    # 标准化数据
                    data_matrix = np.array([[point[obj] for obj in objective_names] for point in pareto_data])
                    normalized_data = (data_matrix - data_matrix.min(axis=0)) / (data_matrix.max(axis=0) - data_matrix.min(axis=0))
                    
                    for i, solution in enumerate(normalized_data):
                        ax.plot(range(len(objective_names)), solution, 'o-', alpha=0.7, linewidth=1)
                    
                    ax.set_xticks(range(len(objective_names)))
                    ax.set_xticklabels(objective_names, fontproperties=self.font_manager.get_font_properties())
                    ax.set_ylabel('标准化目标值', fontproperties=self.font_manager.get_font_properties())
                    ax.set_title('多目标帕累托解集', fontproperties=self.font_manager.get_font_properties(self.config.title_font_size))
                    ax.grid(True, alpha=0.3)
                
                # 保存图表
                if save_path:
                    fig.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
                    return save_path
                else:
                    return self._fig_to_base64(fig)
                    
        except Exception as e:
            self.logger.error(f"帕累托前沿图表创建失败: {e}")
            return self._create_error_chart(f"图表创建失败: {str(e)}")
    
    def create_sensitivity_analysis_chart(self, sensitivity_data: Dict[str, Dict[str, List[float]]], 
                                        save_path: Optional[str] = None) -> str:
        """
        创建敏感性分析图表
        
        Args:
            sensitivity_data: 敏感性数据 {parameter: {'changes': [...], 'impacts': [...]}}
            save_path: 保存路径
            
        Returns:
            str: 图表保存路径或base64编码
        """
        try:
            with self.memory_manager.managed_figure(figsize=(12, 8)) as fig:
                n_params = len(sensitivity_data)
                cols = min(3, n_params)
                rows = (n_params + cols - 1) // cols
                
                for i, (param, data) in enumerate(sensitivity_data.items()):
                    ax = fig.add_subplot(rows, cols, i + 1)
                    
                    changes = data.get('changes', [])
                    impacts = data.get('impacts', [])
                    
                    ax.plot(changes, impacts, 'o-', linewidth=2, markersize=6)
                    ax.axhline(y=0, color='gray', linestyle='--', alpha=0.5)
                    ax.axvline(x=0, color='gray', linestyle='--', alpha=0.5)
                    
                    ax.set_xlabel(f'{param}变化率', fontproperties=self.font_manager.get_font_properties())
                    ax.set_ylabel('目标函数影响', fontproperties=self.font_manager.get_font_properties())
                    ax.set_title(f'{param}敏感性', fontproperties=self.font_manager.get_font_properties())
                    ax.grid(True, alpha=0.3)
                
                fig.suptitle('参数敏感性分析', fontproperties=self.font_manager.get_font_properties(self.config.title_font_size))
                
                # 保存图表
                if save_path:
                    fig.savefig(save_path, dpi=self.config.dpi, bbox_inches='tight')
                    return save_path
                else:
                    return self._fig_to_base64(fig)
                    
        except Exception as e:
            self.logger.error(f"敏感性分析图表创建失败: {e}")
            return self._create_error_chart(f"图表创建失败: {str(e)}")
    
    def _fig_to_base64(self, fig) -> str:
        """
        将图形转换为base64编码
        """
        try:
            import io
            import base64
            
            buffer = io.BytesIO()
            fig.savefig(buffer, format='png', dpi=self.config.dpi, bbox_inches='tight')
            buffer.seek(0)
            
            image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
            buffer.close()
            
            return f"data:image/png;base64,{image_base64}"
            
        except Exception as e:
            self.logger.error(f"图形转换base64失败: {e}")
            return ""
    
    def _create_error_chart(self, error_message: str) -> str:
        """
        创建错误提示图表
        """
        try:
            with self.memory_manager.managed_figure(figsize=(8, 4)) as fig:
                ax = fig.add_subplot(111)
                ax.text(0.5, 0.5, f'图表生成失败\n{error_message}', 
                       ha='center', va='center', fontsize=14,
                       fontproperties=self.font_manager.get_font_properties(),
                       bbox=dict(boxstyle='round', facecolor='lightcoral', alpha=0.8))
                ax.set_xlim(0, 1)
                ax.set_ylim(0, 1)
                ax.axis('off')
                
                return self._fig_to_base64(fig)
                
        except Exception:
            return "图表生成失败"
    
    def get_system_status(self) -> Dict[str, Any]:
        """
        获取系统状态
        """
        return {
            'font_status': {
                'current_font': self.font_manager.current_font,
                'chinese_support': self.font_manager.test_chinese_display()
            },
            'memory_status': self.memory_manager.get_memory_stats(),
            'config': {
                'figure_size': self.config.figure_size,
                'dpi': self.config.dpi,
                'max_memory_mb': self.config.max_memory_mb,
                'max_figures': self.config.max_figures
            }
        }
    
    def cleanup(self):
        """
        清理资源
        """
        try:
            self.memory_manager.cleanup_old_figures(force=True)
            plt.close('all')
            gc.collect()
            self.logger.info("可视化引擎资源清理完成")
        except Exception as e:
            self.logger.error(f"资源清理失败: {e}")


def create_enhanced_visualization_engine(config: Optional[VisualizationConfig] = None) -> EnhancedVisualizationEngine:
    """
    创建增强版可视化引擎实例
    
    Args:
        config: 可视化配置
        
    Returns:
        EnhancedVisualizationEngine: 可视化引擎实例
    """
    return EnhancedVisualizationEngine(config)


# 便捷函数
def quick_scenario_chart(scenario_data: Dict[str, Any], save_path: Optional[str] = None) -> str:
    """
    快速创建情景对比图表
    """
    engine = create_enhanced_visualization_engine()
    return engine.create_scenario_comparison_chart(scenario_data, save_path)


def quick_convergence_chart(convergence_data: Dict[str, List[float]], save_path: Optional[str] = None) -> str:
    """
    快速创建收敛图表
    """
    engine = create_enhanced_visualization_engine()
    return engine.create_optimization_convergence_chart(convergence_data, save_path)


def quick_pareto_chart(pareto_data: List[Dict[str, float]], objective_names: List[str], 
                      save_path: Optional[str] = None) -> str:
    """
    快速创建帕累托前沿图表
    """
    engine = create_enhanced_visualization_engine()
    return engine.create_pareto_frontier_chart(pareto_data, objective_names, save_path)